import ujson as json
import urandom as random
from math import log2


def random_sample(population, k):
    if k > len(population):
        raise ValueError("k cannot be greater than the length of the population")
    population_max_length = len(population) - 1
    return [population[random.randint(0, population_max_length)] for _ in range(0, k)]


class DecisionTree(object):
    def __init__(self, max_depth=5):
        self.max_depth = max_depth
        self.tree = None  # 树的结构

    def fit(self, X, y):
        self.tree = self._build_tree(X, y, depth=0)

    def _build_tree(self, X, y, depth):
        if depth >= self.max_depth or len(set(y)) == 1:
            return max(set(y), key=y.count)  # 返回当前样本中最多的类别

        n_features = len(X[0])
        best_feature, best_threshold = self._best_split(X, y, n_features)

        if best_feature is None:
            return max(set(y), key=y.count)

        left_X, left_y, right_X, right_y = self._split(X, y, best_feature, best_threshold)

        if not left_X or not right_X:
            return max(set(y), key=y.count)

        node = {
            'feature': best_feature,
            'threshold': best_threshold,
            'left': self._build_tree(left_X, left_y, depth + 1),
            'right': self._build_tree(right_X, right_y, depth + 1)
        }
        return node

    def _best_split(self, X, y, n_features):
        best_gini = float('inf')
        best_feature, best_threshold = None, None

        features = random_sample(range(n_features), int(log2(n_features)) + 1)  # 随机选择特征

        for feature in features:
            thresholds = set([x[feature] for x in X])
            for threshold in thresholds:
                left_y, right_y = self._split_labels(X, y, feature, threshold)
                gini = self._gini_index(left_y, right_y)
                if gini < best_gini:
                    best_gini = gini
                    best_feature = feature
                    best_threshold = threshold
        return best_feature, best_threshold

    def _split(self, X, y, feature, threshold):
        left_X, left_y, right_X, right_y = [], [], [], []
        for i in range(len(X)):
            if X[i][feature] <= threshold:
                left_X.append(X[i])
                left_y.append(y[i])
            else:
                right_X.append(X[i])
                right_y.append(y[i])
        return left_X, left_y, right_X, right_y

    def _split_labels(self, X, y, feature, threshold):
        left_y, right_y = [], []
        for i in range(len(X)):
            if X[i][feature] <= threshold:
                left_y.append(y[i])
            else:
                right_y.append(y[i])
        return left_y, right_y

    def _gini_index(self, left_y, right_y):
        n = len(left_y) + len(right_y)
        gini = 0
        for group in [left_y, right_y]:
            if len(group) == 0:
                continue
            p = [group.count(c) / len(group) for c in set(group)]
            gini += (1 - sum([x**2 for x in p])) * (len(group) / n)
        return gini

    def predict(self, x):
        return self._predict_tree(x, self.tree)

    def _predict_tree(self, x, node):
        if isinstance(node, dict):
            if x[node['feature']] <= node['threshold']:
                return self._predict_tree(x, node['left'])
            else:
                return self._predict_tree(x, node['right'])
        else:
            return node


class RandomForest(object):
    def __init__(self, n_trees=10, max_depth=5):
        self.n_trees = n_trees
        self.max_depth = max_depth
        self.trees = []

    def fit(self, X, y):
        for _ in range(self.n_trees):
            tree = DecisionTree(max_depth=self.max_depth)
            sample_X, sample_y = self._bootstrap_sample(X, y)
            tree.fit(sample_X, sample_y)
            self.trees.append(tree)

    def _bootstrap_sample(self, X, y):
        n_samples = len(X)
        sample_X, sample_y = [], []
        for _ in range(n_samples):
            idx = random.randint(0, n_samples - 1)
            sample_X.append(X[idx])
            sample_y.append(y[idx])
        return sample_X, sample_y

    def predict(self, x):
        predictions = [tree.predict(x) for tree in self.trees]
        return max(set(predictions), key=predictions.count)

    def save_model(self, filename):
        """
        保存随机森林模型到文件。

        :param filename: 保存的文件名。
        """
        model_data = []
        for tree in self.trees:
            model_data.append(tree.tree)  # 保存每棵树的结构
        with open(filename, "w") as f:
            json.dump(model_data, f)

    @staticmethod
    def load_model(filename, n_trees=10, max_depth=5):
        """
        从文件加载随机森林模型。

        :param filename: 模型文件名。
        :param n_trees: 树的数量（需与保存时一致）。
        :param max_depth: 树的最大深度（需与保存时一致）。
        :return: 加载的随机森林模型。
        """
        with open(filename, "r") as f:
            model_data = json.load(f)
        
        rf_model = RandomForest(n_trees=n_trees, max_depth=max_depth)
        rf_model.trees = []
        for tree_data in model_data:
            tree = DecisionTree(max_depth=max_depth)
            tree.tree = tree_data
            rf_model.trees.append(tree)
        return rf_model


if __name__ == "__main__":
    if 0:
        # 示例数据集
        X = [
            [2.771244718, 1.784783929],
            [1.728571309, 1.169761413],
            [3.678319846, 2.81281357],
            [3.961043357, 2.61995032],
            [2.999208922, 2.209014212],
            [7.497545867, 3.162953546],
            [9.00220326, 3.339047188],
            [7.444542326, 0.476683375],
            [10.12493903, 3.234550982],
            [6.642287351, 3.319983761]
        ]
        y = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]
        # 创建随机森林
        rf = RandomForest(n_trees=10, max_depth=3)
        rf.fit(X, y)
        # 保存模型
        rf.save_model("/usr/random_forest_model.json")
    else:
        # 加载模型
        loaded_rf = RandomForest.load_model("/usr/random_forest_model.json", n_trees=10, max_depth=3)
        # 使用加载的模型进行预测
        sample = [6.642287351, 3.319983761]
        prediction = loaded_rf.predict(sample)
        print("预测结果:", prediction)
